package main

import (
	"encoding/json"
	"github.com/google/uuid"
	"github.com/pebbe/zmq4"
	"log"
	"time"
)

type TraceData struct {
	Event         string          `json:"event"`
	ProjectID     string          `json:"project_id"`
	WorkflowRunID string          `json:"workflow_run_id"`
	WorkflowID    string          `json:"workflow_id"`
	NodeRunID     string          `json:"node_run_id"`
	NodeID        string          `json:"node_id"`
	Type          int8            `json:"type"`
	Data          json.RawMessage `json:"data"`
}

func main() {
	socket, err := zmq4.NewSocket(zmq4.PUSH)
	if err != nil {
		log.Fatalf("Failed to create socket: %v", err)
	}
	defer socket.Close()

	if err := socket.Connect("tcp://localhost:5555"); err != nil {
		log.Fatalf("Failed to connect: %v", err)
	}

	// 创建示例数据
	projectID := uuid.New().String()
	workflowRunID := uuid.New().String()
	workflowID := uuid.New().String()
	nodeRunID := uuid.New().String()
	nodeID := uuid.New().String()
	userID := uuid.New().String()

	// 发送工作流开始事件
	startData := map[string]interface{}{
		"inputs": map[string]interface{}{
			"keyword":             "test",
			"url":                 "https://example.com",
			"sys.files":           []string{},
			"sys.user_id":         userID,
			"sys.project_id":      projectID,
			"sys.workflow_id":     workflowID,
			"sys.workflow_run_id": workflowRunID,
		},
		"created_at": time.Now().Unix(),
	}
	startDataBytes, _ := json.Marshal(startData)

	workflowStartTrace := TraceData{
		Event:         "workflow_started",
		ProjectID:     projectID,
		WorkflowRunID: workflowRunID,
		WorkflowID:    workflowID,
		NodeRunID:     workflowRunID,
		NodeID:        workflowID,
		Type:          0,
		Data:          startDataBytes,
	}

	// 发送节点开始事件
	nodeStartData := map[string]interface{}{
		"node_type":  "start",
		"title":      "开始",
		"index":      1,
		"inputs":     nil,
		"created_at": time.Now().Unix(),
		"extras":     map[string]interface{}{},
	}
	nodeStartDataBytes, _ := json.Marshal(nodeStartData)

	nodeStartTrace := TraceData{
		Event:         "node_started",
		ProjectID:     projectID,
		WorkflowRunID: workflowRunID,
		WorkflowID:    workflowID,
		NodeRunID:     nodeRunID,
		NodeID:        nodeID,
		Type:          1,
		Data:          nodeStartDataBytes,
	}

	// 发送节点结束事件
	nodeFinishData := map[string]interface{}{
		"node_type":    "start",
		"title":        "开始",
		"index":        1,
		"inputs":       startData["inputs"],
		"outputs":      startData["inputs"],
		"status":       "succeeded",
		"error":        nil,
		"elapsed_time": 0.015796,
		"created_at":   time.Now().Unix(),
		"finished_at":  time.Now().Unix(),
		"files":        []string{},
	}
	nodeFinishDataBytes, _ := json.Marshal(nodeFinishData)

	nodeFinishTrace := TraceData{
		Event:         "node_finished",
		ProjectID:     projectID,
		WorkflowRunID: workflowRunID,
		WorkflowID:    workflowID,
		NodeRunID:     nodeRunID,
		NodeID:        nodeID,
		Type:          1,
		Data:          nodeFinishDataBytes,
	}

	// 发送工作流结束事件
	workflowFinishData := map[string]interface{}{
		"sequence_number": 6,
		"status":          "succeeded",
		"outputs":         nil,
		"error":           nil,
		"elapsed_time":    233.3306336659989,
		"total_tokens":    16117,
		"total_steps":     6,
		"created_by": map[string]interface{}{
			"id":    userID,
			"name":  "test_user",
			"email": "test@example.com",
		},
		"created_at":       time.Now().Unix(),
		"finished_at":      time.Now().Unix(),
		"exceptions_count": 0,
		"files":            []string{},
	}
	workflowFinishDataBytes, _ := json.Marshal(workflowFinishData)

	workflowFinishTrace := TraceData{
		Event:         "workflow_finished",
		ProjectID:     projectID,
		WorkflowRunID: workflowRunID,
		WorkflowID:    workflowID,
		NodeRunID:     workflowRunID,
		NodeID:        workflowID,
		Type:          0,
		Data:          workflowFinishDataBytes,
	}

	// 按顺序发送所有事件
	traces := []TraceData{
		workflowStartTrace,
		nodeStartTrace,
		nodeFinishTrace,
		workflowFinishTrace,
	}

	for _, trace := range traces {
		data, err := json.Marshal(trace)
		if err != nil {
			log.Printf("Failed to marshal trace data: %v", err)
			continue
		}

		if _, err := socket.SendBytes(data, 0); err != nil {
			log.Printf("Failed to send trace data: %v", err)
			continue
		}

		log.Printf("Sent %s event: %s", trace.Event, trace.WorkflowRunID)
		time.Sleep(time.Second) // 等待消息发送完成
	}
}
